# train_test_series.py
import os
import sys
import argparse
import numpy as np
import matplotlib.pyplot as plt
from utils import setup_environment, load_data, save_metrics
import torch
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm

current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if current_dir not in sys.path:
    sys.path.append(current_dir)


from student.time_series import LSTMPredictor

def train(model, train_loader, val_loader, optimizer, scheduler, criterion, 
          device, num_epochs, data_manager, save_dir='checkpoints'):
    os.makedirs(save_dir, exist_ok=True)
    
    train_losses = []
    val_losses = []
    best_val_loss = float('inf')
    best_state_dict = None
    best_epoch = 0
    
    progress_bar = tqdm(total=num_epochs)
    
    try:
        for epoch in range(num_epochs):
            # 训练阶段
            model.train()
            epoch_loss = 0
            for inputs, labels in train_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)
    
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                epoch_loss += loss.item()
    
            epoch_loss /= len(train_loader)
            train_losses.append(epoch_loss)
    
            # 验证阶段
            model.eval()
            val_loss = 0
            with torch.no_grad():
                for inputs, labels in val_loader:
                    inputs = inputs.to(device)
                    labels = labels.to(device)
                    outputs = model(inputs)
                    val_loss += criterion(outputs, labels).item()
            
            val_loss /= len(val_loader)
            val_losses.append(val_loss)
            
            # 学习率调整
            scheduler.step(val_loss)
    
            # 保存最佳模型
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                best_state_dict = model.state_dict()
                best_epoch = epoch
    
            # 更新进度条
            current_lr = optimizer.param_groups[0]['lr']
            progress_bar.set_description(
                f"Epoch {epoch+1}, Train Loss: {epoch_loss:.4f}, Val Loss: {val_loss:.4f}, LR: {current_lr:.2e}"
            )
            progress_bar.update(1)
    
            # 每100个epoch保存checkpoint
            if epoch % 100 == 0:
                checkpoint_path = os.path.join(save_dir, f'checkpoint_epoch_{epoch}.pth')
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'train_loss': epoch_loss,
                    'val_loss': val_loss,
                    'best_val_loss': best_val_loss,
                    'best_epoch': best_epoch
                }, checkpoint_path)
                print(f"\nCheckpoint saved: {checkpoint_path}")
    
            # 打印GPU内存使用
            if torch.cuda.is_available() and epoch % 10 == 0:
                print(f"\nGPU内存使用: {torch.cuda.memory_allocated(0)/1024**2:.1f} MB")
    
    except KeyboardInterrupt:
        print("\n训练被手动中断")
    
    finally:
        progress_bar.close()
        print(f"最佳Epoch: {best_epoch}; 最佳验证损失: {best_val_loss:.4f}")
        
        if best_state_dict is not None:
            model.load_state_dict(best_state_dict)
            save_path = data_manager.save_model(model, mode='serial')
            print(f"最佳模型已保存至: {save_path}")
    
        # 绘制损失曲线
        plt.figure(figsize=(10, 5))
        plt.plot(train_losses, label='训练损失')
        plt.plot(val_losses, label='验证损失')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('训练和验证损失曲线')
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(save_dir, 'training_loss.png'))
        plt.close()

def test(model, test_loader, criterion, device, save_dir='results'):
    os.makedirs(save_dir, exist_ok=True)
    
    model.eval()
    test_loss = 0
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in tqdm(test_loader, desc="测试进度"):
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            test_loss += loss.item()
            all_preds.append(outputs.cpu().numpy())
            all_labels.append(labels.cpu().numpy())
    
    test_loss /= len(test_loader)
    print(f"测试损失: {test_loss:.6f}")
    
    all_preds = np.concatenate(all_preds, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)
    
    # 保存结果
    np.save(os.path.join(save_dir, 'test_predictions.npy'), all_preds)
    np.save(os.path.join(save_dir, 'test_labels.npy'), all_labels)
    
    # 可视化结果
    plt.figure(figsize=(15, 5))
    
    # 1. 时序预测对比图
    plt.subplot(131)
    sample_idx = 0
    port_idx = 0
    plt.plot(all_preds[sample_idx, :, port_idx], 'b-', label='预测值', alpha=0.7)
    plt.plot(all_labels[sample_idx, :, port_idx], 'r-', label='真实值', alpha=0.7)
    plt.title(f'样本 {sample_idx}, 端口 {port_idx} 的时序预测')
    plt.xlabel('时间步')
    plt.ylabel('强度')
    plt.legend()
    plt.grid(True)
    
    # 2. 预测值与真实值的散点图
    plt.subplot(132)
    plt.scatter(all_labels.flatten(), all_preds.flatten(), alpha=0.01)
    plt.plot([all_labels.min(), all_labels.max()], 
            [all_labels.min(), all_labels.max()], 'r--')
    plt.title('预测值 vs 真实值')
    plt.xlabel('真实值')
    plt.ylabel('预测值')
    plt.grid(True)
    
    # 3. 每个端口的误差箱型图
    plt.subplot(133)
    errors = all_preds - all_labels
    plt.boxplot([errors[:,:,i].flatten() for i in range(errors.shape[2])])
    plt.title('各端口预测误差分布')
    plt.xlabel('端口')
    plt.ylabel('误差')
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'test_results.png'))
    plt.close()
    
    # 计算性能指标
    results = {}
    
    # 整体性能指标
    mse = np.mean((all_preds - all_labels) ** 2)
    mae = np.mean(np.abs(all_preds - all_labels))
    rmse = np.sqrt(mse)
    r2 = 1 - np.sum((all_preds - all_labels) ** 2) / np.sum((all_labels - np.mean(all_labels)) ** 2)
    
    results['overall'] = {
        'mse': mse,
        'rmse': rmse,
        'mae': mae,
        'r2': r2
    }
    
    # 各端口性能指标
    results['ports'] = []
    for port in range(all_preds.shape[2]):
        port_preds = all_preds[:,:,port]
        port_labels = all_labels[:,:,port]
        port_mse = np.mean((port_preds - port_labels) ** 2)
        port_mae = np.mean(np.abs(port_preds - port_labels))
        port_rmse = np.sqrt(port_mse)
        port_r2 = 1 - np.sum((port_preds - port_labels) ** 2) / \
                  np.sum((port_labels - np.mean(port_labels)) ** 2)
        
        results['ports'].append({
            'port': port,
            'mse': port_mse,
            'rmse': port_rmse,
            'mae': port_mae,
            'r2': port_r2
        })
    
    save_metrics(results, os.path.join(save_dir, 'test_metrics.npy'))

def main():
    parser = argparse.ArgumentParser(description='训练和测试时序模型')
    parser.add_argument(
        '--mode', type=str, choices=['train', 'test', 'both'], 
        default='both', help='运行模式'
    )
    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')
    parser.add_argument('--epochs', type=int, default=100000, help='训练轮数')
    parser.add_argument('--lr', type=float, default=1e-3, help='学习率')
    args = parser.parse_args()
    
    # 设置环境
    device = setup_environment()
    
    # 加载半径矩阵
    radius_save_path = "../data/radius_matrix.pth"
    radius_matrix = torch.load(radius_save_path, weights_only=False)
    print("半径矩阵加载完成")
    
    # 加载数据
    data_manager, train_dataset, val_dataset, test_dataset = load_data(radius_matrix, mode='serial')
    
    # 创建模型
    model = LSTMPredictor(
        input_dim=11,  # 10个端口 + 1个零通道
        hidden_dim=256,
        output_dim=10,
        num_layers=2,
        dropout=0.1
    ).to(device)
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型总参数量: {total_params:,}")
    
    if args.mode in ['train', 'both']:
        # 训练模型
        train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
        
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
        scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)
        criterion = torch.nn.MSELoss()
        
        train(
            model, train_loader, val_loader, 
            optimizer, scheduler, criterion, 
            device, args.epochs, data_manager
        )
    
    if args.mode in ['test', 'both']:
        # 加载最佳模型
        model = data_manager.load_model(model, mode='serial')
        print("模型加载完成")
        
        # 测试模型
        test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
        criterion = torch.nn.MSELoss()
        
        test(model, test_loader, criterion, device)

if __name__ == "__main__":
    main()